package ai.example.langchain4j.service.impl;

import ai.example.langchain4j.service.LangChain4jService;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentParser;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.RetrievalAugmentor;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.rag.query.router.QueryRouter;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.MemoryId;
import dev.langchain4j.service.UserMessage;
import dev.langchain4j.service.spring.AiService;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingStore;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;

import java.io.IOException;
import java.net.URISyntaxException;
import java.net.URL;
import java.nio.file.FileSystems;
import java.nio.file.Path;
import java.nio.file.PathMatcher;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;

import static dev.langchain4j.data.document.Metadata.metadata;
import static dev.langchain4j.data.document.loader.FileSystemDocumentLoader.loadDocuments;

@Service
@RequiredArgsConstructor
public class LangChain4jServiceImpl implements LangChain4jService {

  private final EmbeddingModel embeddingModel;
  private final ChatModel chatModel;
  private final EmbeddingStore<TextSegment> embeddingStore;
  private final ChatMemoryProvider chatMemoryProvider;

  /**
   * 上传文件并生成向量
   */
  @Override
  public String upload(MultipartFile file) throws IOException {
    if (file == null || file.isEmpty()) {
      return "No file provided";
    }
    String content = new String(file.getBytes());
    TextSegment segment = new TextSegment(content, metadata("fileName", file.getOriginalFilename()));

    embeddingStore.add(embeddingModel.embed(segment).content(), segment);

    return "Upload and embed: " + file.getOriginalFilename();
  }

  /**
   * 搜索向量库
   */
  @Override
  public List<EmbeddingMatch<TextSegment>> search(String message) {

    EmbeddingSearchRequest request = EmbeddingSearchRequest.builder()
      .queryEmbedding(embeddingModel.embed(message).content())
      .maxResults(5)
      .minScore(0.6)
      .build();

    List<EmbeddingMatch<TextSegment>> relevant = embeddingStore.search(request).matches();

    relevant.forEach(m -> {
      System.out.println("Embedding: " + m.embedding());
      System.out.println("Score: " + m.score());
      System.out.println("Text: " + m.embedded().text());
    });

    return relevant;
  }

  /**
   * 创建基于 RAG 的 Assistant
   */
  public AssistantRag createAssistant() {

    // 嵌入 classpath 文档
    EmbeddingStore<TextSegment> docStore = embed(toPath("documents/"), embeddingModel);

    ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
      .embeddingStore(docStore)
      .embeddingModel(embeddingModel)
      .maxResults(5)
      .minScore(0.6)
      .build();

    QueryRouter queryRouter = new QueryRouter() {

      private final String PROMPT_TEMPLATE =
        "Is the following query related to the business of the car rental company? " +
          "Answer only 'yes', 'no' or 'maybe'. Query: {{it}}";

      @Override
      public Collection<ContentRetriever> route(Query query) {
        String message = chatModel.chat(PROMPT_TEMPLATE.replace("{{it}}", query.text()));
        System.out.println("LLM decided: " + message);

        if (message.toLowerCase().contains("no")) {
          return Collections.emptyList();
        }
        return Collections.singletonList(contentRetriever);
      }
    };

    RetrievalAugmentor retrievalAugmentor = DefaultRetrievalAugmentor.builder()
      .queryRouter(queryRouter)
      .build();

    return AiServices.builder(AssistantRag.class)
      .chatModel(chatModel)
      .retrievalAugmentor(retrievalAugmentor)
      .chatMemoryProvider(chatMemoryProvider)
      .build();
  }

  /**
   * 文档嵌入
   */
  private EmbeddingStore<TextSegment> embed(Path documentPath, EmbeddingModel embeddingModel) {

    DocumentParser documentParser = new TextDocumentParser();
    DocumentSplitter splitter = DocumentSplitters.recursive(300, 0);

    List<Document> documents = loadDocuments(documentPath, glob("*.txt"));

    List<TextSegment> segments = new ArrayList<>();
    for (Document doc : documents) {
      segments.addAll(splitter.split(doc));
    }

    List<Embedding> embeddings = embeddingModel.embedAll(segments).content();
    embeddingStore.addAll(embeddings, segments);

    return embeddingStore;
  }

  @AiService
  public interface AssistantRag {
    String chat(@MemoryId String memoryId, @UserMessage String userMessage);
  }

  /**
   * 工具方法
   */
  public static PathMatcher glob(String glob) {
    return FileSystems.getDefault().getPathMatcher("glob:" + glob);
  }

  public static Path toPath(String relativePath) {
    try {
      URL fileUrl = LangChain4jServiceImpl.class.getResource(relativePath);
      return Paths.get(fileUrl.toURI());
    } catch (URISyntaxException e) {
      throw new RuntimeException(e);
    }
  }
}
